import tensorflow.image as timg
import tensorflow as tf
from matplotlib import pyplot as plt
import os
tf.enable_eager_execution()


def random_cut_train(image_path, label_path, train_output_dir = "./data/train/",cut_shape = (256,256),per_cut = 3,prefix = "jpgpng"):
    '''
    将raw图片随机切分为相应大小的n张图片，确保图片和标签文件夹下所有对应的文件名完全相同
    :param imagepath:
    :param labelpath:
    :param image_output_dir:
    :param label_output_dir:
    :return:
    '''

    image_fs = os.listdir(image_path)
    label_fs = os.listdir(label_path)
    image_fs = [os.path.join(image_path,i) for i in image_fs]
    label_fs = [os.path.join(label_path,i) for i in label_fs]

    train_image_dir = os.path.join(train_output_dir,"image")
    train_label_dir = os.path.join(train_output_dir,"label")
    os.makedirs(train_image_dir,exist_ok=True)
    os.makedirs(train_label_dir,exist_ok=True)


    i = 0
    for imagepath,labelpath in zip(image_fs,label_fs):

        # ipath, iname = os.path.split(imagepath)
        # lpath, lname = os.path.split(labelpath)

        # rname, ext = os.path.splitext(iname)
        image, label = tf.read_file(imagepath), tf.read_file(labelpath)
        image, label = timg.decode_image(image), timg.decode_image(label)

        if image.ndim == 3 and image.shape[-1] == 4:
            image = image[:, :, :3]

        if image.shape[:2] != label.shape[:2]:
            print(f"shape not eq in {imagepath}{labelpath}")
            return

        concat = tf.concat([image, label], axis=-1)

        for j in range(per_cut):
            nvalue = timg.random_crop(concat, (cut_shape[0], cut_shape[1], concat.shape[-1]), seed=666)

            clipimage, cliplabel = nvalue[:, :, :image.shape[-1]].numpy(), nvalue[:, :, image.shape[-1]:].numpy()

            ###一个预处理
            cliplabel = cliplabel[:, :, 1]
            cliplabel[cliplabel < 125] = 0
            cliplabel[cliplabel > 125] = 255
            ###


            import scipy.misc
            plt.imsave(train_image_dir, clipimage)
            scipy.misc.toimage(cliplabel).save(train_label_dir + f"/{i}.jpg")
            i+=1


def random_cut_test(image_path,test_dir = "./data/test/", cut_shape = (256,256),per_cut = 2,prefix = ".jpg.png"):
    '''
    将raw图片随机切分为相应的图片
    :param image_path:
    :param test_dir:
    :param cut_shape:
    :param per_cut:
    :param prefix:
    :return:
    '''
    img_fs = os.listdir(image_path)
    img_fs = [os.path.join(image_path,i) for i in img_fs]

    os.makedirs(test_dir,exist_ok=True)

    i = 0
    for imagepath in img_fs:
        ipath, iname = os.path.split(imagepath)
        _, ext = os.path.splitext(iname)# 可能需要判断后缀名

        if ext.lower() not in prefix:
            # print(ext)
            continue

        imagef = tf.read_file(imagepath)
        image = timg.decode_image(imagef)
        if image.ndim == 3 and image.shape[-1] == 4:
            image = image[:, :, :3]

        for j in range(per_cut):
            nvalue = timg.random_crop(image, (cut_shape[0], cut_shape[1], image.shape[-1]), seed=666)
            plt.imsave(os.path.join(test_dir,f"{i}.jpg"), nvalue.numpy())
            i+=1



if __name__ == "__main__":
    random_cut_test(image_path="./rawdata/rawtest/")